use async_trait::async_trait;
use etcd_client::{Client as EtcdClient, PutOptions, LeaseGrantOptions, GetOptions, DeleteOptions, WatchOptions};
use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime};
use tokio::sync::mpsc;
use tokio::time;
use tracing::{debug, error, info, warn};

use super::membership::{Member, MemberStatus, MembershipEvent};
use super::provider::{ClusterConfig, ClusterError, ClusterProvider};

/// etcd-based cluster provider
pub struct EtcdProvider {
    /// etcd client
    client: Arc<RwLock<EtcdClient>>,
    /// etcd configuration
    etcd_config: EtcdConfig,
    /// The ID of the current node
    self_id: RwLock<Option<String>>,
    /// Active lease ID
    lease_id: RwLock<Option<i64>>,
    /// Channel for membership events
    event_tx: mpsc::Sender<MembershipEvent>,
    /// Receive side 
    event_rx: RwLock<Option<mpsc::Receiver<MembershipEvent>>>,
    /// In-memory cache of members
    members_cache: Arc<RwLock<HashMap<String, Member>>>,
}

#[derive(Debug, Clone)]
pub struct EtcdConfig {
    /// etcd endpoints
    pub endpoints: Vec<String>,
    /// Key prefix for cluster nodes
    pub key_prefix: String,
    /// Lease time to live in seconds
    pub lease_ttl: i64,
}

impl Default for EtcdConfig {
    fn default() -> Self {
        Self {
            endpoints: vec!["http://127.0.0.1:2379".to_string()],
            key_prefix: "/protoactor/cluster/".to_string(),
            lease_ttl: 30,
        }
    }
}

impl EtcdProvider {
    /// Create a new etcd provider
    pub async fn new(config: EtcdConfig) -> Result<Self, ClusterError> {
        // Create etcd client
        let client = EtcdClient::connect(config.endpoints.clone(), None)
            .await
            .map_err(|e| ClusterError::ProviderError(format!("Failed to connect to etcd: {}", e)))?;
        
        let (tx, rx) = mpsc::channel(100);
        
        Ok(Self {
            client: Arc::new(RwLock::new(client)),
            etcd_config: config,
            self_id: RwLock::new(None),
            lease_id: RwLock::new(None),
            event_tx: tx,
            event_rx: RwLock::new(Some(rx)),
            members_cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }
    
    /// Get key for a cluster node
    fn get_node_key(&self, cluster_name: &str, host: &str, port: u16) -> String {
        format!("{}{}/{}/{}", self.etcd_config.key_prefix, cluster_name, host, port)
    }
    
    /// Create a lease for automatic member removal
    async fn create_lease(&self) -> Result<i64, ClusterError> {
        let mut client = {
            let client_guard = self.client.read().unwrap();
            client_guard.clone()
        };
        
        let options = LeaseGrantOptions::new().with_ttl(self.etcd_config.lease_ttl);
        let resp = client.lease_grant(self.etcd_config.lease_ttl, Some(options))
            .await
            .map_err(|e| ClusterError::ProviderError(format!("Failed to create lease: {}", e)))?;
        
        let lease_id = resp.id();
        
        // Store lease ID
        {
            let mut lease_id_guard = self.lease_id.write().unwrap();
            *lease_id_guard = Some(lease_id);
        }
        
        Ok(lease_id)
    }
    
    /// Keep the lease alive
    async fn keep_lease_alive(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
        let client = self.client.clone();
        
        tokio::spawn(async move {
            let lease_id = {
                let lease_id_guard = self.lease_id.read().unwrap();
                match *lease_id_guard {
                    Some(id) => id,
                    None => {
                        error!("Cannot keep lease alive - no lease ID");
                        return;
                    }
                }
            };
            
            // Keep lease alive
            loop {
                let mut client_instance = {
                    let client_guard = client.read().unwrap();
                    client_guard.clone()
                };
                
                match client_instance.lease_keep_alive(lease_id).await {
                    Ok(_) => {
                        debug!("Successfully refreshed lease");
                    }
                    Err(e) => {
                        error!("Failed to refresh lease: {}", e);
                        // TODO: Attempt to create a new lease
                    }
                }
                
                // Sleep for a duration less than the lease TTL
                let sleep_duration = Duration::from_secs((self.etcd_config.lease_ttl as u64) / 3);
                time::sleep(sleep_duration).await;
            }
        })
    }
    
    /// Convert etcd key-value to Member
    fn kv_to_member(&self, key: &str, value: &str) -> Option<Member> {
        // Extract ID from the key
        // Expected key format: /protoactor/cluster/cluster_name/host/port
        let parts: Vec<&str> = key.split('/').collect();
        if parts.len() < 6 {
            warn!("Invalid key format: {}", key);
            return None;
        }
        
        let cluster_name = parts[3];
        let host = parts[4];
        let port_str = parts[5];
        
        let port = match port_str.parse::<u16>() {
            Ok(p) => p,
            Err(_) => {
                warn!("Invalid port in key: {}", key);
                return None;
            }
        };
        
        // Parse the JSON value
        let member_data: Result<serde_json::Value, _> = serde_json::from_str(value);
        let member_data = match member_data {
            Ok(data) => data,
            Err(e) => {
                warn!("Invalid member data JSON: {}", e);
                return None;
            }
        };
        
        // Create new member
        let mut member = Member::new(
            key.to_string(),
            host.to_string(),
            port,
        );
        
        // Extract role if available
        if let Some(role) = member_data.get("role").and_then(|r| r.as_str()) {
            member = member.with_role(role);
        }
        
        // Extract metadata if available
        if let Some(metadata) = member_data.get("metadata").and_then(|m| m.as_object()) {
            for (key, value) in metadata {
                if let Some(value_str) = value.as_str() {
                    member.metadata.insert(key.clone(), value_str.to_string());
                }
            }
        }
        
        // Set status and last heartbeat
        member.status = MemberStatus::Alive;
        member.last_heartbeat = SystemTime::now();
        
        Some(member)
    }
    
    /// Start watching for member changes
    pub fn start_watcher(
        self: Arc<Self>,
        config: ClusterConfig
    ) -> tokio::task::JoinHandle<()> {
        let prefix_key = format!("{}{}", self.etcd_config.key_prefix, config.name);
        let client = self.client.clone();
        let members_cache = self.members_cache.clone();
        let event_tx = self.event_tx.clone();
        
        tokio::spawn(async move {
            loop {
                let mut client_instance = {
                    let client_guard = client.read().unwrap();
                    client_guard.clone()
                };
                
                // Set up watch with prefix
                let options = WatchOptions::new().with_prefix();
                
                match client_instance.watch(prefix_key.clone(), Some(options)).await {
                    Ok(mut watcher) => {
                        info!("Started watching etcd for member changes at prefix: {}", prefix_key);
                        
                        // Process watch events
                        while let Some(resp) = watcher.message().await.map_err(|e| {
                            error!("Watch error: {}", e);
                        }).ok().flatten() {
                            for event in resp.events() {
                                match event.event_type() {
                                    etcd_client::EventType::Put => {
                                        if let Some(kv) = event.kv() {
                                            let key = match kv.key_str() {
                                                Ok(k) => k.to_string(),
                                                Err(_) => {
                                                    warn!("Invalid UTF-8 in key");
                                                    continue;
                                                }
                                            };
                                            
                                            let value = match kv.value_str() {
                                                Ok(v) => v.to_string(),
                                                Err(_) => {
                                                    warn!("Invalid UTF-8 in value");
                                                    continue;
                                                }
                                            };
                                            
                                            // Convert to member
                                            if let Some(member) = self.kv_to_member(&key, &value) {
                                                // Check if this is a new member
                                                let is_new = {
                                                    let cache = members_cache.read().unwrap();
                                                    !cache.contains_key(&key)
                                                };
                                                
                                                // Update cache
                                                {
                                                    let mut cache = members_cache.write().unwrap();
                                                    cache.insert(key.clone(), member.clone());
                                                }
                                                
                                                // Send joined event if new
                                                if is_new {
                                                    let event = MembershipEvent::MemberJoined(super::membership::MemberJoined {
                                                        member: member.clone(),
                                                    });
                                                    let _ = event_tx.try_send(event);
                                                    info!("Member joined: {}", key);
                                                }
                                            }
                                        }
                                    },
                                    etcd_client::EventType::Delete => {
                                        if let Some(kv) = event.kv() {
                                            let key = match kv.key_str() {
                                                Ok(k) => k.to_string(),
                                                Err(_) => {
                                                    warn!("Invalid UTF-8 in key");
                                                    continue;
                                                }
                                            };
                                            
                                            // Get the member and mark as down
                                            let member = {
                                                let mut cache = members_cache.write().unwrap();
                                                if let Some(member) = cache.get_mut(&key) {
                                                    member.status = MemberStatus::Down;
                                                    member.clone()
                                                } else {
                                                    // Member not in cache
                                                    continue;
                                                }
                                            };
                                            
                                            // Send down event
                                            let event = MembershipEvent::MemberDown(super::membership::MemberDown {
                                                member: member.clone(),
                                            });
                                            let _ = event_tx.try_send(event);
                                            info!("Member down: {}", key);
                                        }
                                    }
                                }
                            }
                        }
                    },
                    Err(e) => {
                        error!("Failed to watch etcd: {}", e);
                    }
                }
                
                // Wait before retrying
                time::sleep(Duration::from_secs(5)).await;
            }
        })
    }
    
    /// Load all existing members from etcd
    async fn load_members(&self, cluster_name: &str) -> Result<(), ClusterError> {
        let mut client = {
            let client_guard = self.client.read().unwrap();
            client_guard.clone()
        };
        
        let prefix_key = format!("{}{}", self.etcd_config.key_prefix, cluster_name);
        let options = GetOptions::new().with_prefix();
        
        let resp = client.get(prefix_key.clone(), Some(options))
            .await
            .map_err(|e| ClusterError::ProviderError(format!("Failed to get members from etcd: {}", e)))?;
        
        // Process all key-values
        for kv in resp.kvs() {
            let key = match kv.key_str() {
                Ok(k) => k.to_string(),
                Err(_) => {
                    warn!("Invalid UTF-8 in key");
                    continue;
                }
            };
            
            let value = match kv.value_str() {
                Ok(v) => v.to_string(),
                Err(_) => {
                    warn!("Invalid UTF-8 in value");
                    continue;
                }
            };
            
            // Convert to member
            if let Some(member) = self.kv_to_member(&key, &value) {
                // Store in cache
                {
                    let mut cache = self.members_cache.write().unwrap();
                    cache.insert(key.clone(), member.clone());
                }
                
                // We don't emit events here as this is initial load
                debug!("Loaded member from etcd: {}", key);
            }
        }
        
        Ok(())
    }
}

#[async_trait]
impl ClusterProvider for EtcdProvider {
    async fn register(&self, config: &ClusterConfig) -> Result<(), ClusterError> {
        // Create a lease for this member
        let lease_id = self.create_lease().await?;
        
        // Generate node key
        let key = self.get_node_key(&config.name, &config.host, config.port);
        
        // Build metadata JSON
        let mut metadata = HashMap::new();
        metadata.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string());
        metadata.insert("system_type".to_string(), "protoactor-rs".to_string());
        
        let member_data = serde_json::json!({
            "role": config.name,
            "metadata": metadata
        });
        
        // Register with etcd
        let mut client = {
            let client_guard = self.client.read().unwrap();
            client_guard.clone()
        };
        
        let options = PutOptions::new().with_lease(lease_id);
        client.put(key.clone(), member_data.to_string(), Some(options))
            .await
            .map_err(|e| ClusterError::ProviderError(format!("Failed to register with etcd: {}", e)))?;
        
        // Store self ID
        {
            let mut self_id = self.self_id.write().unwrap();
            *self_id = Some(key.clone());
        }
        
        // Load existing members
        self.load_members(&config.name).await?;
        
        // Create local member record
        let mut member = Member::new(key.clone(), config.host.clone(), config.port)
            .with_role(&config.name);
        
        member.metadata.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string());
        member.metadata.insert("system_type".to_string(), "protoactor-rs".to_string());
        member.status = MemberStatus::Alive;
        member.last_heartbeat = SystemTime::now();
        
        // Store in cache
        {
            let mut cache = self.members_cache.write().unwrap();
            cache.insert(key.clone(), member.clone());
        }
        
        // Emit member joined event for ourselves
        let event = MembershipEvent::MemberJoined(super::membership::MemberJoined {
            member: member.clone(),
        });
        
        let _ = self.event_tx.try_send(event);
        
        info!("Node registered with etcd, key: {}", key);
        Ok(())
    }
    
    async fn deregister(&self) -> Result<(), ClusterError> {
        // Get our ID
        let key = {
            let self_id = self.self_id.read().unwrap();
            match &*self_id {
                Some(id) => id.clone(),
                None => return Err(ClusterError::ProviderError("Node not registered".to_string())),
            }
        };
        
        // Delete from etcd
        let mut client = {
            let client_guard = self.client.read().unwrap();
            client_guard.clone()
        };
        
        client.delete(key.clone(), None)
            .await
            .map_err(|e| ClusterError::ProviderError(format!("Failed to deregister from etcd: {}", e)))?;
        
        // Get our member record and update status
        let member = {
            let mut cache = self.members_cache.write().unwrap();
            
            match cache.get_mut(&key) {
                Some(member) => {
                    member.status = MemberStatus::Left;
                    member.clone()
                },
                None => return Err(ClusterError::MemberNotFound(key)),
            }
        };
        
        // Emit member left event
        let event = MembershipEvent::MemberLeft(super::membership::MemberLeft {
            member,
        });
        
        let _ = self.event_tx.try_send(event);
        
        info!("Node deregistered from etcd: {}", key);
        Ok(())
    }
    
    async fn get_members(&self) -> Result<Vec<Member>, ClusterError> {
        let cache = self.members_cache.read().unwrap();
        let members = cache.values().cloned().collect();
        Ok(members)
    }
    
    async fn monitor_members(&self) -> mpsc::Receiver<MembershipEvent> {
        let mut rx_opt = self.event_rx.write().unwrap();
        match rx_opt.take() {
            Some(rx) => rx,
            None => {
                // Create a new channel if the first one was already taken
                let (tx, rx) = mpsc::channel(100);
                self.event_tx.clone();
                rx
            }
        }
    }
    
    fn as_any(&self) -> &dyn Any {
        self
    }
} 